#include "StdAfx.h"
#include <QFileDialog>
#include <QMessageBox>
#include <QTextStream>
#include <QString>
#include <QProgressDialog>
#include <iostream>
using namespace std;
#include "mainwindow.h"

#include "volumereader.h"
#include "CAttribute.h"
#include "Attrs.h"
#include "scoped_array.h"
#include "qglobal.h"

static QTextStream os(stdout);

QTextStream & operator<<(QTextStream & os,CMFT* mft) {

	if (mft->haveTimes())
		os<<mft->getPath()+mft->getFileName()<<",";
	else return os;
	NTFSFileTime times = mft->getTimes();
	os<<ToString(&times.fileNameTime.entry)<<",";
	os<<ToString(&times.fileNameTime.access)<<",";
	os<<ToString(&times.fileNameTime.creation)<<",";
	os<<ToString(&times.fileNameTime.modify)<<",";
	os<<ToString(&times.stdInfoTime.entry)<<",";
	os<<ToString(&times.stdInfoTime.access)<<",";
	os<<ToString(&times.stdInfoTime.creation)<<",";
	os<<ToString(&times.stdInfoTime.modify)<<",";
	return os;
}

enum COLUMNS {
	COL_FILE_NAME=0,
	COL_START_BYTE,
	COL_IS_DIR,
	COL_MFT_NO,
	COL_SIZE,
	COL_FNM,
	COL_FNA,
	COL_FNC,
	COL_FNE,
	COL_SIM,
	COL_SIA,
	COL_SIC,
	COL_SIE,
	COL_IS_DEL,
	COL_COLLISION=COL_START_BYTE
};

void fillItem( QTreeWidgetItem * nxt, CMFT* cur ) 
{
	nxt->setData(0,Qt::UserRole,cur->diskOffSet);
	if (cur->countAttr(AttrType::FileName)>=1)
		nxt->setData(COL_FILE_NAME,Qt::DisplayRole,cur->getPath()+cur->getFileName());
	nxt->setData(COL_START_BYTE,Qt::DisplayRole,cur->diskOffSet);
	nxt->setData(COL_MFT_NO,Qt::DisplayRole,cur->entryNo());
	nxt->setData(COL_IS_DIR,Qt::DisplayRole,cur->isDirectory());
	if (cur->haveTimes()) {
		nxt->setData(COL_FNM,Qt::DisplayRole,ToString(&(cur->getTimes().fileNameTime.modify)));
		nxt->setData(COL_FNA,Qt::DisplayRole,ToString(&(cur->getTimes().fileNameTime.access)));
		nxt->setData(COL_FNC,Qt::DisplayRole,ToString(&(cur->getTimes().fileNameTime.creation)));
		nxt->setData(COL_FNE,Qt::DisplayRole,ToString(&(cur->getTimes().fileNameTime.entry)));
		nxt->setData(COL_SIM,Qt::DisplayRole,ToString(&(cur->getTimes().stdInfoTime.modify)));
		nxt->setData(COL_SIA,Qt::DisplayRole,ToString(&(cur->getTimes().stdInfoTime.access)));
		nxt->setData(COL_SIC,Qt::DisplayRole,ToString(&(cur->getTimes().stdInfoTime.creation)));
		nxt->setData(COL_SIE,Qt::DisplayRole,ToString(&(cur->getTimes().stdInfoTime.entry)));
	}
	nxt->setData(COL_IS_DEL,Qt::DisplayRole,cur->isDeleted());
	//TimeBasedInference(nxt,cur);
}

MainWindow::MainWindow(QWidget *parent)
	: QMainWindow(parent)
{
	ui.setupUi(this);
	ui.dateTimeStart->setDateTime( QDateTime::currentDateTime().addMonths(-1));
	ui.dateTimeEnd->setDateTime(QDateTime::currentDateTime());
	enableControls(false);
	setupConnections();
	//init tree widget;
	ui.treeDeletedFiles->setSortingEnabled(true);
	ui.treeOrphanFiles->setSortingEnabled(true);
	QTreeWidgetItem* header = new QTreeWidgetItem();
	header->setData(COL_FILE_NAME,Qt::DisplayRole,"FileName");
	header->setData(COL_MFT_NO,Qt::DisplayRole,"MFT NO.");
	header->setData(COL_START_BYTE,Qt::DisplayRole,"StartByte");
	header->setData(COL_FNM,Qt::DisplayRole,"FN M");
	header->setData(COL_FNA,Qt::DisplayRole,"FN A");
	header->setData(COL_FNC,Qt::DisplayRole,"FN C");
	header->setData(COL_FNE,Qt::DisplayRole,"FN E");
	header->setData(COL_SIM,Qt::DisplayRole,"SI M");
	header->setData(COL_SIA,Qt::DisplayRole,"SI A");
	header->setData(COL_SIC,Qt::DisplayRole,"SI C");
	header->setData(COL_SIE,Qt::DisplayRole,"SI E");
	header->setData(COL_SIZE,Qt::DisplayRole,"Size");
	header->setData(COL_IS_DEL,Qt::DisplayRole,"Deleted");
	header->setData(COL_IS_DIR,Qt::DisplayRole,"Directory");
	ui.treeDeletedFiles->setHeaderItem(header);
	QTreeWidgetItem *orphan = new QTreeWidgetItem(*header);
	ui.treeOrphanFiles->setHeaderItem(orphan);
#ifdef _DEBUG
	OpenImage("\\\\.\\i:");
#endif
}


MainWindow::~MainWindow()
{

}

void MainWindow::ExportFile()
{
	QTextStream os(stdout);
	QTextStream is(stdin);
	QList<QTreeWidgetItem*> selected = ui.treeDeletedFiles->selectedItems();
	if (selected.size()!=1) {
		selected = ui.treeOrphanFiles->selectedItems();
	}
	if (selected.size()==1) {
		QTreeWidgetItem* current = selected.at(0);
		UINT64 Diskoffset = (current->data(0,Qt::UserRole)).toLongLong();
		CMFT mft = vr->MFTAt(Diskoffset);
		printMFTInfo(&mft);
		os<<"Input Attr Type:";
		os.flush();
		int typeId;
		is>>typeId;
		os<<"Input ID:";
		os.flush();
		int dataId;
		is>>dataId;
		CAttribute dataAtt = mft.getAttribute(typeId,dataId);

		QString fileName = QFileDialog::getSaveFileName(0,"Recover as",mft.getFileName());
		//really write the file
		if (!fileName.isNull()) {
			QFile outFile(fileName);
			outFile.open(QFile::WriteOnly);
			QDataStream dataStream(&outFile);
			scoped_array<char> buffer(new char[vr->clusterSize()]);
			UINT64 i;
			for ( i = 0;i<dataAtt.size()/vr->clusterSize();i++) {
				dataAtt.read(i*vr->clusterSize(),vr->clusterSize(),buffer.get());
				dataStream.writeRawData(buffer.get(),vr->clusterSize());
			}
			if (dataAtt.size()%vr->clusterSize()!=0) { //< write the tails
				UINT64 start = i*vr->clusterSize();
				dataAtt.read(start,dataAtt.size()-start,buffer.get());
				dataStream.writeRawData(buffer.get(),dataAtt.size()-start);
			}
		}
	}
}

void MainWindow::setupConnections()
{
	connect(ui.action_Open,SIGNAL(triggered()),this,SLOT(OpenImage()));
	connect(ui.action_Export_File,SIGNAL(triggered()),this,SLOT(ExportFile()));
	connect(ui.action_Export_Log,SIGNAL(triggered()),this,SLOT(ExportLog()));
}

void MainWindow::enableControls( bool enable )
{
	// TODO: Add standard code header comment here.
}

void MainWindow::OpenImage(QString imagePath)
{
	QDateTime dt = ui.dateTimeStart->dateTime();
	startTime = ToSYSTEMTIME(dt);
	endTime = ToSYSTEMTIME(ui.dateTimeEnd->dateTime());
	QFileDialog fileDialog(this);
	fileDialog.setFileMode(QFileDialog::AnyFile);
	if (imagePath=="")
		if (fileDialog.exec()) {
			QStringList fileName = fileDialog.selectedFiles();
			if (fileName.size()!=1) {
				QMessageBox::warning(this,"you did not select one file","NTReco");
				return;
			}
			imagePath = fileName[0];
		}
		vr.reset(new VolumeReader(imagePath));
		if (!vr->isReady()) {
			QMessageBox::warning(this,vr->getError(),"NTReco");
		} else {
			//scan&store all mft 
			refillTreeWidget();
			//calculate mft collision
			calcCollision();
			//display timeline
			refillTimeLine();
		}

}
void MainWindow::ExportLog() {
	QString logName = QFileDialog::getSaveFileName();
	if (logName.isNull())
		return;
	if (timeLine.writeLog(logName)) {
		QMessageBox::information(this,"OK","OK");
	}else QMessageBox::information(this,"Error writing log",logName);

}
void MainWindow::refillTreeWidget() {
	//clear ui
	ui.treeDeletedFiles->clear();
	ui.treeOrphanFiles->clear();
	mftMap.clear();
	QTextStream os(stdout);
	QTextStream is(stdin);
	//fill
	CMFT masterMFT = vr->MFTEntry(0);
	Q_ASSERT(masterMFT.countAttr(AttrType::Data)==1);
	//xprintMFTInfo(&masterMFT);
	//xstd::cout<<"input $DATA id:";
	int masterDataId;
	for (int i = 0;i<masterMFT.nresAttrs.size();i++)
		if (masterMFT.nresAttrs[i]->Header.TypeID==AttrType::Data)
			masterDataId = masterMFT.nresAttrs[i]->Header.ID;
	//xis>>masterDataId;
	CAttribute masterData = masterMFT.getAttribute(AttrType::Data,masterDataId);
	char buffer[1024];
	QProgressDialog progress("Reading MFT...", "Abort", 0, masterData.size(), this);
	progress.setWindowModality(Qt::WindowModal);
	//scan all mft and store in mftMap
	for (UINT64 i = 0;i < masterData.size();i+=1024) {
		//progress
		if (i%(1024*100)==0)
			progress.setValue(i);
		if (progress.wasCanceled())
			break;
		//work
		UINT64 diskOffset = masterMFT.getAttPos(AttrType::Data,masterDataId,i);
		if(!vr->isMFT(diskOffset)) {
			std::cout<<"Not MFT AT entry :"<<i/1024<<"\n";
			continue;
		}
		CMFT cur= vr->MFTAt(diskOffset);
		//printMFTInfo(&cur);
		if (!cur.hasExisted())
			//exclude mft that have never been used
			continue;
		if (!cur.isBase()) {
			os<<"NONE BASE MFT AT"<<i<<endl;
			continue;
		}
		//if reach here then |cur| is a valid mft
		this->mftMap[diskOffset] = cur;
	}
	//finish progressdialog
	progress.setValue(masterData.size());
}

void MainWindow::calcCollision()
{
	QFile csvFile("report.csv");	
	csvFile.open(QIODevice::WriteOnly);
	QTextStream csv(&csvFile);
	QProgressDialog progress("Calculate Collisions...","Abort",0,mftMap.size(),this);
	progress.setModal(true);
	csv<<"FILE NAME,MFT M,A,C,E,STDINFO M,A,C,E,COLLISION FILES(DISKOFFSET-FILENAME-COLLISION"<<endl;
	int nProgress = 0;
	timeLine.clear();
	for (QMap<UINT64,CMFT>::iterator itMe = mftMap.begin();!(itMe==mftMap.end());itMe++,nProgress++) {
		//progress
		if (nProgress%10==0) {
			progress.setValue(nProgress);
			if (progress.wasCanceled())
				return;
		}
		CMFT * me = &itMe.value();
		if (me->hasExisted()/*DEBUG:*/ ) {//< only deleted could have collision
			csv<<me;//output current file
			QTreeWidgetItem *nxt = new QTreeWidgetItem();
			//
			// all info including timebased inference
			fillItem(nxt,me);
			//BUILD TIMELINE
			timeLine.TimeBasedInference(nxt ,me,startTime,endTime);
			bool isOrphan = me->isOrphan();
			if (isOrphan) 
				ui.treeOrphanFiles->addTopLevelItem(nxt);
			else 
				ui.treeDeletedFiles->addTopLevelItem(nxt);
			//
			//collision
// 			for (QMap<UINT64,CMFT>::iterator itOther = mftMap.begin();itOther!=mftMap.end();itOther++) {
// 				CMFT * other = &itOther.value();
// 				if (other->entryNo()==8) //< other is $BadCluster
// 					continue;
// 				if (me==other) 
// 					continue;//< self collision
// 				QList<Range> collision  = me->getCollision(other);
// 				if (collision.size()>0) {
// 					QTreeWidgetItem * ncol = new QTreeWidgetItem();
// 					//
// 					// all info including timebase inference
// 					fillItem(ncol,other);
// 					nxt->addChild(ncol);
// 					csv<<other->diskOffSet;
// 					csv<<","<<other->getPath()+other->getFileName();
// 					QString colPos = "";
// 					for (int i = 0;i<collision.size();i++) {
// 						csv<<"("<<collision[i].start<<'-'<<collision[i].end<<")";
// 						colPos += QString("(%1,%2)").arg(collision[i].start).arg(collision[i].end);
// 					}
// 					QTreeWidgetItem *colString = new QTreeWidgetItem();
// 					colString->setData(COL_COLLISION,Qt::DisplayRole,colPos);
// 					ncol->addChild(colString);
// 					os<<endl;
// 				}
// 			}
			csv<<endl;
		}

	}
	csv.flush();
	progress.setValue(progress.maximum());
}

void MainWindow::refillTimeLine()
{
	timeLine.fillTree(ui.treeTimeLine);
}

 